#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
@Time    : 2025/4/3 16:56
@Author  : HZP
@File    : sqlalchemy.py
"""
from contextlib import contextmanager

from flask_sqlalchemy import SQLAlchemy as _SQLAlchemy
from injector import inject


@inject
class SQLAlchemy(_SQLAlchemy):
    """重写Flask-SQLAlchemy中的核心类，实现自动提交"""

    @contextmanager
    def auto_commit(self):
        # 定义一个上下文管理器，用于自动提交和回滚事务
        try:
            # 尝试执行yield语句
            yield
            # 如果没有异常，则提交事务
            self.session.commit()
        except Exception as e:
            # 如果有异常，则回滚事务，并抛出异常
            self.session.rollback()
            print(e)
            raise e
